#include "proc.h"
#include "panic.h"
#include "page.h"
#include "riscv.h"
#include "pmm.h"
#include "fs.h"
#include "elf.h"
#include "list.h"
#include "sbi.h"
#include "string.h"
#include "app.h"
#include "stdio.h"

extern struct scheduler sche_fifo;
struct scheduler *sche_ptr = &sche_fifo;

// only the positive value is valid
static int pid = 1;
// init proc
struct proc *init_proc;
struct context init_context;
// pointer of proc being running
struct proc *cur_proc;
// process pool
struct proc proc_pool[MAX_PROC_NUM];

/* @brief: initialize proc_pool and init_proc
 */
void proc_init() {
	init_proc = &proc_pool[0];
	init_proc->state = RUNNING;
	init_proc->pid = pid++;
	init_proc->con_ptr = &init_context;
	init_proc->table_ptr3 = &kpgtable;
	init_fd_table(init_proc->fd_table);
	cur_proc = init_proc;
	int i;
	for (i = 1; i < MAX_PROC_NUM; ++i) {
		proc_pool[i].pid = 0;
		proc_pool[i].state = UNINIT;
		proc_pool[i].tf_ptr = NULL;
		proc_pool[i].con_ptr = NULL;
	}
	sche_ptr->init();
}

/* @brief: add a process to runnable queue
 */
void run_proc(struct proc *proc_ptr) {
	sche_ptr->runnable(proc_ptr);
}

/* @brief: alloc an unused struct proc
 */
struct proc *alloc_proc() {
	int i;
	for (i = 1; i < MAX_PROC_NUM; ++i) {
		if (proc_pool[i].state == UNINIT) {
			proc_pool[i].pid = pid++;
			return &proc_pool[i];
		}
	}
	panic("Too many process.");
	return NULL;
}

/* @brief: alloc
 */
int alloc_fd() {
	int i = 4;
	while (cur_proc->fd_table[i] != NULL) {
		i++;
	}
	return i;
}

/* @brief: initialize proc's fd_table
 * 	stdin, stdout, stderr, root directory
 */
void init_fd_table(struct file_info *fd_table[]) {
	int i;
	fd_table[stdin] = &file_table[stdin];
	fd_table[stdout] = &file_table[stdout];
	fd_table[stderr] = &file_table[stderr];
	fd_table[ROOT_DIR_FD] = &file_table[ROOT_DIR_FD];
	for (i = 4; i < PROC_FD_NUM; ++i) {
		fd_table[i] = NULL;
	}
}

/* @brief: free fd_table
 * @param proc: fd_table to be free
 */
void free_fd_table(struct file_info *fd_table[]) {
	int i;
	for (i = 4; i < PROC_FD_NUM; ++i) {
		if (fd_table[i] != NULL) {
			fd_table[i]->type = UNUSED;
		}
	}
}

/* @brief: free struct proc's vma list
 * @param vma_ptr: start pointer of the vma list
 */
void free_vma_list(struct vma** vma_dptr) {
	struct vma *p;
	while ((*vma_dptr) != NULL) {
		p = (struct vma*)list_rm_front((struct linked_list **)vma_dptr);
		kfree((uint64)p, sizeof(struct vma));
	}
}

/* @brief: free user app's code
 * @param vma_ptr: starting struct vma
 */
void free_segments(struct vma *vma_ptr) {
	while (vma_ptr != NULL) {
		kfree(vma_ptr->start_addr, vma_ptr->size);
		vma_ptr = (struct vma*)vma_ptr->next;
	}
}

/* @brief: copy user app's code to a new address and map to a page table
 * @param vma_ptr: indicates the valid virtual address
 * @param table_ptr3: map new physical address to this level3 page table
 * 	the virtual address are identical with vma_ptr
 */
void copy_segments(struct vma *vma_ptr, struct pgtable_l3 *table_ptr3) {
	uint64 loaded;
	uint64 total_size;
	uint64 size;
	uint64 pa;
	uint64 va;
	uint64 offset;
	while (vma_ptr != NULL) {
		loaded = 0;
		total_size = vma_ptr->size;
		va = vma_ptr->start_addr;
		// if va is not PGSIZE aligned, fill 0 for the start of 1st page
		if ((offset = (va - ROUNDDOWN(va, PGSIZE))) != 0) {
			pa = kalloc(PGSIZE);
			size = PGSIZE - offset;
			memset((char *)pa, 0, offset);
			memcpy((char *)(pa + offset), (char *)va, size);
			map_a_page(table_ptr3, ROUNDDOWN(va, PGSIZE), pa, vma_ptr->flags);
			loaded += size;
			va += size;
		}
		// copy and map the rest pages
		while (loaded < total_size) {
			pa = kalloc(PGSIZE);
			size = total_size - loaded > PGSIZE ? PGSIZE : total_size - loaded;
			memcpy((char *)pa, (char *)va, size);
			map_a_page(table_ptr3, ROUNDDOWN(va, PGSIZE), pa, vma_ptr->flags);
			loaded += size;
			va += size;
		}
		vma_ptr = (struct vma *)vma_ptr->next;
	}
}

/* @brief: copy struct trapframe
 * @param to:
 * @param from:
 */
void copy_trapframe(struct trapframe *to, struct trapframe *from) {
	memcpy((void *)to, (void *)from, sizeof(struct trapframe));
}

/* @brief: proc schedule and create new user app
 * @note: when the current last process call this, a new user app will be created
 */
void proc_switch() {
	struct proc *last_ptr = cur_proc;
	while ((cur_proc = sche_ptr->next_proc()) == NULL) {
		cur_proc = init_proc;
		csrw(satp, (SV39 << 60) | ((uint64)cur_proc->table_ptr3 >> 12));
		asm volatile ("sfence.vma");
		proc_wait(-1, NULL, 0);
		app_init();
	}
	// change proc state
	last_ptr->state = RUNNABLE;
	cur_proc->state = RUNNING;
	// change page table
	csrw(satp, (SV39 << 60) | ((uint64)cur_proc->table_ptr3 >> 12));
	asm volatile ("sfence.vma");
	context_switch(last_ptr->con_ptr, cur_proc->con_ptr);
}

/* @brief: clone a process
 * @param flags
 * @param stack: os allocated when stack is 0
 * 	parent assigned when stack is not 0
 * @param ptid, tls, ctid: not used now
 * @return: 0 implies this is child process, -1 failed
 * 	>0 implies this is parent process, ret is the child process id
 * @note: parent proc will continue to run after clone, child proc in the runnable queue
 */
int proc_clone(uint64 flags, uint64 stack, uint64 ptid, uint64 tls, uint64 ctid) {
	struct trapframe *tf_ptr;
	struct context *con_ptr;
	struct vma *vma_ptr = cur_proc->vma_list;
	uint64 ustack;
	uint64 user_kstack;
	struct proc *proc_ptr;
	struct pgtable_l3 *table_ptr3;

	// create a new page table and map kernel codes
	table_ptr3 = (struct pgtable_l3 *)kalloc(sizeof(struct pgtable_l3));
	smode_map_pages(table_ptr3);
	// copy code and map pages
	copy_segments(vma_ptr, table_ptr3);

	// parent proc assigned user stack
	if (stack != 0) {
		ustack = stack;
	}
	// os allocated user stack
	else {
		ustack = kalloc(USTACK_SIZE);
		// copy user stack
		memcpy((void *)ustack, (void *)ROUNDDOWN(cur_proc->tf_ptr->sp, USTACK_SIZE), USTACK_SIZE);
		ustack += (cur_proc->tf_ptr->sp) - ROUNDDOWN(cur_proc->tf_ptr->sp, USTACK_SIZE);
	}
	// allocate user's kernel stack
	user_kstack = kalloc(PGSIZE);

	// create a new trapframe
	tf_ptr = (struct trapframe *)kalloc(sizeof(struct trapframe));
	copy_trapframe(tf_ptr, cur_proc->tf_ptr);
	// user stack top
	tf_ptr->sp = (uint64)ustack;
	// kernel stack top
	tf_ptr->kstack = (uint64)user_kstack + PGSIZE;
	// child process return 0
	tf_ptr->a0 = 0;

	// create a new context
	con_ptr = (struct context *)kalloc(sizeof(struct context));
	con_ptr->ra = (uint64)trap_ret_user;
	con_ptr->sp = (uint64)tf_ptr->kstack;

	// allocate a struct proc
	proc_ptr = alloc_proc();
	proc_ptr->tf_ptr = tf_ptr;
	proc_ptr->con_ptr = con_ptr;
	proc_ptr->state = RUNNABLE;
	proc_ptr->table_ptr3 = table_ptr3;
	proc_ptr->vma_list = vma_ptr;
	proc_ptr->vma_list->count++;
	proc_ptr->parent_ptr = cur_proc;
	proc_ptr->zombie_child_list = NULL;

	// init fd_table
	init_fd_table(proc_ptr->fd_table);

	// add child process to the runnable queue
	run_proc(proc_ptr);
	// parent child: return child process's pid
	return proc_ptr->pid;
}

/* @brief: execute an user app
 * @param path: file path
 * @param argv: arguments
 * @param envp: environment variables
 * @return: -1 failed, 0 successful
 */
int proc_execve(char *path, char * argv[], char *envp[]) {
	uint64 file_ptr;
	uint64 fsize;
	uint64 app_entry;
	int fd;
	
	// open new file
	if ((fd = fs_openat(AT_FDCWD, path, O_RDWR, O_RDWR)) == -1) {
		// file not found
		return -1;
	}
	// free original code segments
	free_segments(cur_proc->vma_list);
	// free original vma_list, add a new one
	if (--(cur_proc->vma_list->count) < 1) {
		free_vma_list(&(cur_proc->vma_list));
	}
	cur_proc->vma_list = NULL;
	fsize = cur_proc->fd_table[fd]->size;
	file_ptr = kalloc(fsize);
	// read file
	if (fs_read(fd, (void *)file_ptr, fsize) == -1) {
		panic("proc_execve: read file failed\n");
	}
	// load elf file
	app_entry = load_elf((Elf64_Ehdr *)(file_ptr), cur_proc->table_ptr3, &(cur_proc->vma_list));
	// free file
	kfree(file_ptr, fsize);
	// update sepc and return address
	// user and kernel stack can keep the previous one
	cur_proc->tf_ptr->epc = app_entry;
	cur_proc->con_ptr->ra = (uint64)trap_ret_user;
	// refresh TLB
	asm volatile ("sfence.vma");
	return 0;
}

/* @brief: free struct proc and its variables
 * @param proc_ptr: pointer to the proc to be free
 */
void free_proc(struct proc *proc_ptr) {
	kfree((uint64)(proc_ptr->con_ptr), sizeof(struct context));
	kfree((uint64)(proc_ptr->tf_ptr), sizeof(struct trapframe));
	if (--(proc_ptr->vma_list->count) < 1) {
		free_vma_list(&(proc_ptr->vma_list));
	}
	proc_ptr->zombie_child_list = NULL;
	proc_ptr->parent_ptr = NULL;
	proc_ptr->state = UNINIT;
	// TODO: free level 1 and level 2 page table
	kfree((uint64)(proc_ptr->table_ptr3), sizeof(struct pgtable_l3));
}


/* @brief: current process wait for an child process to end
 * @param pid: if -1, wait for any child process, else wait for the specific one
 * @param code: child process exit status
 * @param options: not used now
 * @return: -1 failed, otherwise child pid
 */
int proc_wait(int pid, int *code, int options) {
	struct zombie *p;
	int find = 0;
	// wait for the child process whose pid is param pid
	if (pid != -1) {
		while (!find) {
			while ((p = cur_proc->zombie_child_list) == NULL) {
				proc_switch();
			}
			while (p->proc_ptr->pid != pid && p->next != NULL) {
				p = (struct zombie *)p->next;
			}
			if (p->proc_ptr->pid == pid) {
				find = 1;
			}
		}
	}
	// pid = -1, wait for any child process to end
	else {
		while ((p = cur_proc->zombie_child_list) == NULL) {
			proc_switch();
		}
	}
	list_rm((struct linked_list **)(&(cur_proc->zombie_child_list)), (struct linked_list *)p);
	// param code may be NULL
	if (code != NULL) {
		*code = p->exit_val << 8;
	}
	pid = p->proc_ptr->pid;
	free_proc(p->proc_ptr);
	kfree((uint64)p, sizeof(struct zombie));
	return pid;
}

/* @brief: process exit
 * @param ec: exited value
 */
void proc_exit(int ec) {
	struct zombie *zombie_ptr;

	free_segments(cur_proc->vma_list);
	free_fd_table(cur_proc->fd_table);
	// create zombie node for its parent process
	zombie_ptr = (struct zombie *)kalloc(sizeof(struct zombie));
	zombie_ptr->exit_val = ec;
	zombie_ptr->proc_ptr = cur_proc;
	zombie_ptr->next = NULL;
	list_add_front((struct linked_list **)(&(cur_proc->parent_ptr->zombie_child_list)), (struct linked_list *)zombie_ptr);
	// if current process has any zombie process
	while (cur_proc->zombie_child_list != NULL) {
		proc_wait(-1, NULL, 0);
	}
	cur_proc->state = ZOMBIE;
	// switch current process to init_proc
	cur_proc = init_proc;
	// choose another user process
	proc_switch();
}

